#include <algorithm>
#include <vector>

#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/layers/deconv_layer.hpp"

namespace caffe {

template <typename Dtype>
void DeconvolutionLayer<Dtype>::compute_output_shape() {
  const int* kernel_shape_data = this->kernel_shape_.cpu_data();
  const int* stride_data = this->stride_.cpu_data();
  const int* pad_data = this->pad_.cpu_data();
  const int* dilation_data = this->dilation_.cpu_data();
  this->output_shape_.clear();
  for (int i = 0; i < this->num_spatial_axes_; ++i) {
    // i + 1 to skip channel axis
    const int input_dim = this->input_shape(i + 1);
    const int kernel_extent = dilation_data[i] * (kernel_shape_data[i] - 1) + 1;
    const int output_dim = stride_data[i] * (input_dim - 1)
        + kernel_extent - 2 * pad_data[i];
    this->output_shape_.push_back(output_dim);
  }
}

template <typename Dtype>
void DeconvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  // Load asymmetric padding parameters.
  ConvolutionParameter conv_param = this->layer_param_.convolution_param();
  int pad_l_=0, pad_r_=0, pad_t_=0, pad_b_=0;
  pad_l_ = conv_param.pad_left();
  pad_r_ = conv_param.pad_right();
  pad_t_ = conv_param.pad_top();
  pad_b_ = conv_param.pad_bottom();

  // Create temporary Blobs to record forward results.
  vector<Blob<Dtype>*> tmp;
  
  for (int i = 0; i < bottom.size(); ++i) {
    int N, C, H, W, top_h, top_w, offset_h=0, offset_w=0;
    N = top[i]->shape(0);
    C = top[i]->shape(1);
    H = top[i]->shape(2);
    W = top[i]->shape(3);
    top_h = H;
    top_w = W;
    if(pad_l_ != pad_r_ || pad_t_ != pad_b_) {
      top_h -= abs(pad_t_ - pad_b_);
      top_w -= abs(pad_l_ - pad_r_);
      offset_h = pad_t_ > pad_b_ ? (pad_t_ - pad_b_) : 0;
      offset_w = pad_l_ > pad_r_ ? (pad_l_ - pad_r_) : 0;
    }
    // Reshape top Blobs according to asymmetric padding parameters.
    top[i]->Reshape(N, C, top_h, top_w);
    // Create tmp Blobs.
    tmp.push_back(new Blob<Dtype>(N, C, H, W));

    // Forward propagate.
    const Dtype* bottom_data = bottom[i]->cpu_data();
    Dtype* tmp_data = tmp[i]->mutable_cpu_data();
    Dtype* top_data = top[i]->mutable_cpu_data();
    const int tmp_dim_ = tmp[i]->count(1);
    for (int n = 0; n < this->num_; ++n) {
      const Dtype* weight = this->blobs_[0]->cpu_data();
      this->backward_cpu_gemm(bottom_data + n * this->bottom_dim_, weight,
          tmp_data + n * tmp_dim_);
      if (this->bias_term_) {
        const Dtype* bias = this->blobs_[1]->cpu_data();
        this->forward_cpu_bias(tmp_data + n * tmp_dim_, bias);
      }
    }
    
    // Copy data from temporary Blobs to top Blobs.
    for(int n = 0; n < N; ++n) {
      for(int c = 0; c < C; ++c) {
        for(int h = 0; h < top_h; ++h) {
          for(int w = 0; w < top_w; ++w) {
            top_data[((n * C + c) * top_h + h) * top_w + w] = tmp_data[((n * C + c) * H + h + offset_h) * W + w + offset_w];
          }
        }
      }
    }
  }

  // Release Memory of temporary Blobs.
  for (int i = 0; i < bottom.size(); ++i) {
    delete tmp[i];
  }
}

template <typename Dtype>
void DeconvolutionLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  const Dtype* weight = this->blobs_[0]->cpu_data();
  Dtype* weight_diff = this->blobs_[0]->mutable_cpu_diff();
  for (int i = 0; i < top.size(); ++i) {
    const Dtype* top_diff = top[i]->cpu_diff();
    const Dtype* bottom_data = bottom[i]->cpu_data();
    Dtype* bottom_diff = bottom[i]->mutable_cpu_diff();
    // Bias gradient, if necessary.
    if (this->bias_term_ && this->param_propagate_down_[1]) {
      Dtype* bias_diff = this->blobs_[1]->mutable_cpu_diff();
      for (int n = 0; n < this->num_; ++n) {
        this->backward_cpu_bias(bias_diff, top_diff + n * this->top_dim_);
      }
    }
    if (this->param_propagate_down_[0] || propagate_down[i]) {
      for (int n = 0; n < this->num_; ++n) {
        // Gradient w.r.t. weight. Note that we will accumulate diffs.
        if (this->param_propagate_down_[0]) {
          this->weight_cpu_gemm(top_diff + n * this->top_dim_,
              bottom_data + n * this->bottom_dim_, weight_diff);
        }
        // Gradient w.r.t. bottom data, if necessary, reusing the column buffer
        // we might have just computed above.
        if (propagate_down[i]) {
          this->forward_cpu_gemm(top_diff + n * this->top_dim_, weight,
              bottom_diff + n * this->bottom_dim_,
              this->param_propagate_down_[0]);
        }
      }
    }
  }
}

#ifdef CPU_ONLY
STUB_GPU(DeconvolutionLayer);
#endif

INSTANTIATE_CLASS(DeconvolutionLayer);

}  // namespace caffe
